package main

import (
    "os"
    "syscall"
    "unsafe"
    "github.com/gliderlabs/ssh"
    gossh "golang.org/x/crypto/ssh"
    "github.com/kr/pty"
    "io"
    "fmt"
    "strings"
    "os/exec"
)

func setWinsize(f *os.File, w, h int) {//调整term窗口大小使用 不需要动
    syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ), uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}



func (dp *SSHProxy) publicKeyAuth(ctx ssh.Context, key ssh.PublicKey) bool { //使用公钥登陆的验证
    dp.mu.RLock()
    defer dp.mu.RUnlock()
    cli:=ctx.RemoteAddr().String()
    for _,app:=range dp.AppInfos{
        if app.TrustedClientKeys[cli]==fmt.Sprintf("%x",gossh.MarshalAuthorizedKey(key)){
            return true
        }
    }
    return false

}

func (dp *SSHProxy) passwdAuth(ctx ssh.Context, password string) bool {//使用密码登陆的验证
    dp.mu.RLock()
    defer dp.mu.RUnlock()
    app, ok := dp.AppInfos[ctx.User()]
    return ok && app.Auth(password, dp.Salt)
}

func (dp *SSHProxy) serve() error {
    ssh.Handle(dp.handle)
    srv := &ssh.Server{Addr: dp.Port, Handler: nil}
    srv.SetOption(ssh.PasswordAuth(dp.passwdAuth)) //使用密码登陆的设置
    srv.SetOption(ssh.PublicKeyAuth(dp.publicKeyAuth)) //使用密码登陆的设置
    srv.SetOption(ssh.HostKeyFile(dp.PemFilePath))
    return srv.ListenAndServe()
}

func (dp *SSHProxy)AuthContainer(appid,cid string)bool{//验证容器id是否有权限
    dp.mu.RLock()
    defer dp.mu.RUnlock()
    if cid==""{
        return true
    }
    app,ok:=dp.AppInfos[appid]
    if !ok{
        return false
    }
    return app.AuthContainer(cid)
}

func (dp *SSHProxy) handle(session ssh.Session) {//处理
    pkstr:=fmt.Sprintf("%x",gossh.MarshalAuthorizedKey(session.PublicKey()))
    cli:=session.RemoteAddr().String()
    appid := session.User()
    command, cid := ParseCommandAndCid(appid, session.Command())
    if len(command) == 0 {
        io.WriteString(session, fmt.Sprintf("command %v is not supported\n", session.Command()))
        session.Exit(1)
        return
    }
    io.WriteString(session, fmt.Sprintf("command is: [%s]\n", strings.Join(command, " ")))
    if command[0]=="reg"{//reg 命令用来设置免密码ssh 以后屏蔽
        dp.RegClientKey(appid,cli,pkstr)
        io.WriteString(session, "authorize key accepted\n")
        session.Exit(0)
        return
    }
    if cid != "" && !dp.AuthContainer(appid,cid) {
        io.WriteString(session, fmt.Sprintf("%s is not allowed to %s\n", cid, appid))
        session.Exit(1)
        return
    }
    ptyReq, winCh, isPty := session.Pty()
    if !isPty {
        io.WriteString(session, "No PTY requested\n try to use:\n ssh -tt xx.xx.xx.xx command\n")
        session.Exit(1)
        return
    }
    cmd := exec.Command(command[0], command[1:]...)
    cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
    f, err := pty.Start(cmd)
    if err != nil {
        io.WriteString(session, fmt.Sprintf("command error\n %v\n", err))
        session.Exit(1)
    }
    go func() {//调整窗口大小
        for win := range winCh {
            setWinsize(f, win.Width, win.Height)
        }
    }()
    go func() {//读输入
        io.Copy(f, session) // stdin
    }()
    io.Copy(session, f) // stdout
}
